import os
import torch
import cv2
from one_layer_model import IlluNet_with_Quad
import pyiqa
import os
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

def to_deploy(mo):
    for m in mo.modules():
        if hasattr(m, 'switch_to_deploy'):
            m.switch_to_deploy()
    return mo



def run():
    device = torch.device('cuda')
    niqe_metric = pyiqa.create_metric('niqe', device=device).eval()
    PI_metric = pyiqa.create_metric('pi', device=device).eval()


    model = IlluNet_with_Quad(3,3,None).eval()
    model.load_state_dict(torch.load(args.pre_weight))
    model = to_deploy(model).to(device)



    seq_list =  ('LIME','DICM','MEF','NPE','VV')

    for seq in seq_list:
        avg_PI_list = []
        avg_NIQE_list = []
        out_p = os.path.join(args.out_path,seq)
        if not os.path.exists(out_p):
            os.makedirs(out_p)
        seq_p = os.path.join(args.data_path,seq)

        img_list = sorted(os.listdir(seq_p))
        for img_n in img_list:
            img_p = os.path.join(seq_p,img_n)
            img_np = cv2.imread(img_p).astype('float32')/255.0
            img_ten = torch.from_numpy(img_np).permute(2,0,1).unsqueeze(0).to(device)
            with torch.no_grad():
                out = model(img_ten)
                score_niqe = niqe_metric(out.float())
                score_pi = PI_metric(out.float())
                print('{}  PI:  {} NIQE: {} '.format(seq+'/'+img_n,score_pi.item(),score_niqe.item()))
                avg_PI_list.append(score_pi.item())
                avg_NIQE_list.append(score_niqe.item())

                enhanced_np = out.detach().cpu().squeeze(0).permute(1,2,0).numpy()*255.0
                cv2.imwrite(out_p+'/'+img_n.split('.')[0]+'.png',enhanced_np)
        print('sequence is {} avg PI is {} avg NIQE is {}'.format(seq,sum(avg_PI_list)/len(avg_PI_list),sum(avg_NIQE_list)/len(avg_NIQE_list)))


if __name__=='__main__':
    device = torch.device('cuda')   
    parser = ArgumentParser(description="validation script for refence",formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('--pre_weight', default='./weight/sclm_noref.pth', type=str, help='weight of model')
    parser.add_argument('--data_path', default='xxx', type=str, help='input data')
    parser.add_argument('--out_path', default= 'xxx', type=str, help='output path')

    args = parser.parse_args()
    run()